/*
	Universidade de Sao Paulo - 2014-02
	Programacao Concorrente - Turma 12 Grupo A
	---
  Bruno Augusto dos Santos - 7696389
	Franco Lamping - 7696351
	Gabriel Massarin de Paula Aguiar - 7573002
*/

#include <iostream>
#include <fstream>
#include <cstdlib>
#include <string>
#include <ctime>
#include <cmath>
#include <mpi.h>

#define INPUT  "in.ppm"
#define OUTPUT "out.ppm"
#define MASK_SIZE 5
#define MASTER 0

/* Estrutura de pixel com as cores RGB */
typedef struct {
  unsigned char r,g,b;
} pixel;

/*
* Estrutura que define cada secao da imagem
* w e h: largura e altura
* sX, sY: posicao X e Y do inicio da secao
* duration: tempo de processamento
* *image: ponteiro para a estrutura pixel
*/
typedef struct {
  int w, h;
  int sX, sY;
  double duration;
  pixel *image;
} Section;

pixel *in;

/* Le as cores RGB da imagem ppm e as copiam para o vetor global 'in' */
void coletarRGB(int *width, int *height) {

    int depth;
    std::string magicNumber;
    std::ifstream inputFile;
    inputFile.open(INPUT);
    if (inputFile == NULL) {
        std::cout << "Failed to open input file." << std::endl;
        exit(1);
    }

    /* Leitura dos primeiros valores do cabecalho da imagem
    *  magicNumber: assinatura da imagem
    */
    inputFile >> magicNumber >> (*width) >> (*height) >> depth;

    /* Verifica se imagem esta no formato permitido */
    if( magicNumber != "P3" ) {
      std::cout << "Imagem incompativel!" << std::endl;
      exit(1);
    }

    /* Aloca e salva os valores RGB em 'in' */
    in = (pixel*) malloc(sizeof(pixel)*(*width)*(*height));
    int r,g,b;
    for(int i = 0; i < (*height); i++){
        for(int j = 0; j < (*width); j++){
            inputFile >> r >> g >> b;
            in[(*width)*i + j].r = r;
            in[(*width)*i + j].g = g;
            in[(*width)*i + j].b = b;
        }
    }
    inputFile.close();
}

/*
* Aplica o filtro da Media nos valores contidos em 'in' e os salvam em 'out'
* in: valores RGB lidos da imagem
* out: valores RGB com a aplicao do filtro
*/
void applyFilter (pixel *in, int width, int height, pixel *out) {
	for(int i = 0; i < height; i++){
		for(int j = 0; j < width; j++){
			int R = 0, G = 0, B = 0;
			for(int k = -MASK_SIZE/2; k <= MASK_SIZE/2; k++){
				for(int l = -MASK_SIZE/2; l <= MASK_SIZE/2; l++){
					if( (i + k >= 0) && (i + k < height) && (j + l >= 0) && (j + l < width) ){
						R += in[(i + k)*width + (j + l)].r;
						G += in[(i + k)*width + (j + l)].g;
						B += in[(i + k)*width + (j + l)].b;
					}
				}
			}
			out[i*width + j].r = R/(MASK_SIZE*MASK_SIZE);
			out[i*width + j].g = G/(MASK_SIZE*MASK_SIZE);
			out[i*width + j].b = B/(MASK_SIZE*MASK_SIZE);
		}
	}
}

/*
* Imprime os valores contidos em 'out' no arquivo de saida,
* gerando a imagem final com o filtro da Media aplicado
*/
void salvarImagemFinal (int width, int height, int depth, pixel *out) {
    std::ofstream outputFile;
    outputFile.open(OUTPUT);
    outputFile << "P3" << std::endl << width << " " << height << std::endl << depth << std::endl;
    for(int i = 0; i < height; i++){
        for(int j = 0; j < width; j++){
            outputFile << static_cast<unsigned>(out[(width*i) + j].r) << " ";
            outputFile << static_cast<unsigned>(out[(width*i) + j].g) << " ";
            outputFile << static_cast<unsigned>(out[(width*i) + j].b) << " ";
        }
        outputFile << std::endl;
    }
    outputFile.close();
}

int main(int argc, char **argv){
	pixel *out;
	std::string magicNumber;
	int width, height, depth;
	int taskid, numtasks, source, dest, tag = 1;
	MPI_Status status;

  /* Inicia MPI e armazena o id do processo e o numero de processos */
	MPI_Init(&argc,&argv);
	MPI_Comm_rank(MPI_COMM_WORLD,&taskid);
	MPI_Comm_size(MPI_COMM_WORLD,&numtasks);

  /*
  * Cria uma estrutura em MPI do tipo pixel
  * Este procedimento e importante para que possamos passar estrutuas
  * como mensagem para os outros processos
  */
  const int itemsPixel = 3;
  int lengthsPixel[3]   = {1,1,1};
  MPI_Datatype typesPixel[3] = {MPI_UNSIGNED_CHAR, MPI_UNSIGNED_CHAR, MPI_UNSIGNED_CHAR};
  MPI_Aint offsets[3];
  MPI_Datatype mpi_pixel;

  offsets[0] = offsetof(pixel, r);
  offsets[1] = offsetof(pixel, g);
  offsets[2] = offsetof(pixel, b);

  MPI_Type_create_struct(itemsPixel, lengthsPixel, offsets, typesPixel, &mpi_pixel);
  MPI_Type_commit(&mpi_pixel);

  /*
  * Cria uma estrutura em MPI do tipo Section
  * Note que esta estrutura em MPI nao contem a referencia para pixel,
  * uma vez que esta referencia e invalida nos outros nodes
  * Os valores de pixel sao enviados separadamente de Section
  */
  const int itemsSection = 5;
  int lengthsSection[5] = {1, 1, 1, 1, 1};
  MPI_Datatype typesSection[5] = {MPI_INT, MPI_INT, MPI_INT, MPI_INT, MPI_DOUBLE};
  MPI_Datatype mpi_Section;
  MPI_Aint     offsets2[5];

  offsets2[0] = offsetof(Section, w);
  offsets2[1] = offsetof(Section, h);
  offsets2[2] = offsetof(Section, sX);
  offsets2[3] = offsetof(Section, sY);
  offsets2[4] = offsetof(Section, duration);

  MPI_Type_create_struct(itemsSection, lengthsSection, offsets2, typesSection, &mpi_Section);
  MPI_Type_commit(&mpi_Section);

  /**************************** Node Principal********************************/
  if (taskid == MASTER) {
		std::cout << "mpi_mm iniciou com " << numtasks << " tasks." << std::endl;

    int xSections = atoi(argv[1]);
    int ySections = atoi(argv[2]);

    if( xSections * ySections != numtasks ) {
      std::cout << "O produto dos argumentos precisa coincidir com o numero de processos. Ex: " << std::endl;
      std::cout << "mpirun -np P executavel X Y" << std::endl;
      std::cout << "onde P = X*Y" << std::endl;
      exit(1);
    }

    /* Vetor de Section que armazena os valores RGB iniciais de cada secao da imagem ppm */
    Section **sections = (Section **) malloc( ySections * sizeof(Section*) );
    for(int i=0; i<ySections; i++) {
      sections[i] = (Section *) malloc( xSections * sizeof(Section));
    }

    /* Vetor de Section que armazena os valores RGB com o filtro aplicado de cada secao da imagem ppm */
    Section **sectionsOut = (Section **) malloc( ySections * sizeof(Section*) );
    for(int i=0; i<ySections; i++) {
      sectionsOut[i] = (Section *) malloc( xSections * sizeof(Section));
    }

    /* Funcao de leitura da imagem */
    coletarRGB(&width, &height);

    /* Calcula valores relevantes de cada secao da imagem */
    int secWidth  = width/xSections  + MASK_SIZE -1;
    int secHeight = height/ySections + MASK_SIZE -1;
    int realSecWidth  = width/xSections;
    int realSecHeight = height/ySections;

    /* Aloca 'out' levando em conta as cores RGB */
    out = new pixel[width*height];

    /* Gera as secoes da imagem */
    for(int i=0; i<ySections; i++) {
      for(int j=0; j<xSections; j++) {
        sections[i][j].image = new pixel[secWidth*secHeight];
        /* Calcula largura e altura da ultima section para ajustar ao tamanho da imagem */
        sections[i][j].w = (j == xSections-1) ? (width - realSecWidth*j) + MASK_SIZE-1  : secWidth;
        sections[i][j].h = (i == ySections-1) ? (height - realSecHeight*i)  + MASK_SIZE-1  : secHeight;
        sections[i][j].sY = i*realSecHeight-2;
        sections[i][j].sX = j*realSecWidth-2;

        sectionsOut[i][j].image = new pixel[secWidth*secHeight];
        /* Calcula largura e altura da ultima section para ajustar ao tamanho da imagem */
        sectionsOut[i][j].w = (j == xSections-1) ? (width - realSecWidth*j) + MASK_SIZE-1  : secWidth;
        sectionsOut[i][j].h = (i == ySections-1) ? (height - realSecHeight*i)  + MASK_SIZE-1  : secHeight;
        sectionsOut[i][j].sY = i*realSecHeight-2;
        sectionsOut[i][j].sX = j*realSecWidth-2;
      }
    }

    /* Atribui os valores RGB de cada secao */
    for(int i=0; i<ySections; i++) {
        for(int j=0; j<xSections; j++) {

          int si, sj;
          for(int m=sections[i][j].sY, si=0; m<sections[i][j].sY+sections[i][j].h; m++, si++) {
            for(int n=sections[i][j].sX, sj=0; n<sections[i][j].sX+sections[i][j].w; n++, sj++) {
              if (m < 0 || m >= height || n < 0 || n >= width) {
                pixel p;
                p.r = p.g = p.b = 0;
                sections[i][j].image[ si*sections[i][j].w + sj ] = p;
              } else {
                sections[i][j].image[ si*sections[i][j].w + sj ] = in[ m * width + n ];
              }

            }
          }

      }
    }

    /* Envia uma secao da imagem para o respectivo processo */
    for(int i=0; i<ySections; i++) {
      for(int j=0; j<xSections; j++) {
        if(i == 0 && j == 0)
          continue; // Trabalho do processo principal

        /* Envia a Section */
        MPI_Send(&sections[i][j], 1, mpi_Section, i*xSections+j, tag, MPI_COMM_WORLD);

        /* Envia o pixel */
        MPI_Send(sections[i][j].image, sections[i][j].w*sections[i][j].h, mpi_pixel, i*xSections+j, tag+1, MPI_COMM_WORLD);
      }
    }

    /* Registra o inicio do processamento */
    time_t t = clock();

    /* Aplica o filtro na secao*/
    applyFilter(sections[0][0].image, sections[0][0].w, sections[0][0].h, sectionsOut[0][0].image);

    /* Registra o fim do processamento */
    t = clock() - t;
    sections[0][0].duration = ((float)t)/CLOCKS_PER_SEC;

    /* Coleta as secoes de cada processo */
    Section tmp;
    for(int i=0; i<numtasks-1; i++) {
      MPI_Recv(&tmp, 1, mpi_Section, MPI_ANY_SOURCE, tag, MPI_COMM_WORLD, &status);
      int ti = (status.MPI_SOURCE)/xSections;
      int tj = (status.MPI_SOURCE)%xSections;
      sectionsOut[ti][tj].w = tmp.w;
      sectionsOut[ti][tj].h = tmp.h;
      sectionsOut[ti][tj].sX = tmp.sX;
      sectionsOut[ti][tj].sY = tmp.sY;
      sectionsOut[ti][tj].duration = tmp.duration;
      MPI_Recv(sectionsOut[ti][tj].image, tmp.w*tmp.h, mpi_pixel, status.MPI_SOURCE, tag+1, MPI_COMM_WORLD, &status);
    }

    /*
    * Constroi o vetor de pixel 'out' com os valores RGB de cada secao ja
    * com o filtro aplicado
    */
    double max_time = 0.0;
    for(int i=0; i<ySections; i++) {
        for(int j=0; j<xSections; j++) {
          max_time = (sectionsOut[i][j].duration > max_time) ? sectionsOut[i][j].duration : max_time;
          int si, sj;
          for(int m=sectionsOut[i][j].sY+2, si=2; m<sectionsOut[i][j].sY+sectionsOut[i][j].h-2; m++, si++) {
            for(int n=sectionsOut[i][j].sX+2, sj=2; n<sectionsOut[i][j].sX+sectionsOut[i][j].w-2; n++, sj++) {
              if (m < 0 || m >= height || n < 0 || n >= width) {
                continue;
              } else {
                out[ m * width + n ] = sectionsOut[i][j].image[ si*sectionsOut[i][j].w + sj ];
              }
            }
          }
      }
    }

    std::cout << "Elapsed time (maximum among " << numtasks << " processes): " << max_time << " seconds." << std::endl;

    /*
    * Imprime os valores contidos em 'out' no arquivo de saida,
    * gerando a imagem final com o filtro da Media aplicado
    */
    salvarImagemFinal(width, height, 255, out);

    delete[] in;
    delete[] out;

	}
  /**************************** Nodes adicionais ********************************/
  else {

    /* Aguarda o recebimento da Section */
    Section tmpIn, tmpOut;
    MPI_Recv(&tmpIn, 1, mpi_Section, 0, tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

    tmpIn.image = new pixel[tmpIn.w*tmpIn.h];
    tmpOut.image = new pixel[tmpIn.w*tmpIn.h];
    tmpOut.w = tmpIn.w;
    tmpOut.h = tmpIn.h;
    tmpOut.sX = tmpIn.sX;
    tmpOut.sY = tmpIn.sY;

    /* Aguarda o recebimento de pixel */
    MPI_Recv(tmpIn.image, tmpIn.w*tmpIn.h, mpi_pixel, 0, tag+1, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

    /* Registra o inicio do processamento */
    time_t t = clock();

    /* Aplica o filtro na secao */
    applyFilter(tmpIn.image, tmpIn.w, tmpIn.h, tmpOut.image);

    /* Registra o fim do processamento */
    t = clock() - t;
    tmpOut.duration = ((float)t)/CLOCKS_PER_SEC;

    /* Retorna a secao da imagem com o filtro aplicado */
    MPI_Send(&tmpOut, 1, mpi_Section, 0, tag, MPI_COMM_WORLD);
    MPI_Send(tmpOut.image, tmpOut.w*tmpOut.h, mpi_pixel, 0, tag+1, MPI_COMM_WORLD);
  }

  /* Finaliza MPI */
  MPI_Finalize();

  return 0;
}
